import os
import threading
import time

import cv2
import h5py
import numpy as np
from tqdm import tqdm
from sahi.slicing import slice_image
from sahi.utils.cv import read_image
from ultralytics import YOLO
from concurrent.futures import ThreadPoolExecutor, as_completed

# 配置参数
INPUT_VIDEO_PATH = 'E:/EFY/002无人机视频/DJI_0813.MP4'
HDF5_PATH = 'E:/EFY/003无人机数据/DJI_0813_4.h5'
OUTPUT_VIDEO_PATH = 'E:/EFY/003无人机数据/DJI_0813-3.MP4'
MODEL_PATH = 'weights/visdrone+mydata100+crowdhuman200.pt'

SLICE_SIZE = 300
OVERLAP_RATIO = 0.1
CONF_THRESHOLD = 0.6
# 全局状态变量

import cv2
import h5py
import numpy as np
import threading
import time
from queue import Queue
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor
from sahi.slicing import slice_image

# ================= 全局监控变量 ===================
lock = threading.Lock()
total_submitted = 0
total_processed = 0
active_futures = []
monitor_stop = False
frame_queue = Queue(maxsize=64)  # 控制帧缓存大小
write_queue = Queue(maxsize=128)  # 控制写入任务缓存大小

# ================= 监控线程 ===================
def monitor_thread():
    while not monitor_stop:
        with lock:
            print(f"[MONITOR] 已提交: {total_submitted} | 已完成: {total_processed} | 活跃任务: {len(active_futures)} | 缓存帧: {frame_queue.qsize()} | 待写入: {write_queue.qsize()}")
        time.sleep(1)

# ============== 解码线程：从视频读取帧放入缓存 ==============
def decode_thread_func(video_path):
    cap = cv2.VideoCapture(video_path)
    while True:
        if frame_queue.full():
            time.sleep(0.05)
            continue
        ret, frame = cap.read()
        if not ret:
            frame_queue.put(None)  # 结束标志
            break
        resized = cv2.resize(frame, (1280, 720))
        frame_queue.put(resized)
    cap.release()

# ============== 写入线程：从 write_queue 中取出写入 HDF5 ==============
def write_worker(h5_file):
    while True:
        item = write_queue.get()
        if item is None:
            break
        frame_idx, patch_images, patch_origins = item
        try:
            grp = h5_file.create_group(f'frame_{frame_idx}')
            for i, (patch, (ox, oy)) in enumerate(zip(patch_images, patch_origins)):
                try:
                    dset = grp.create_dataset(f'patch_{i}', data=patch, compression='lzf')
                    dset.attrs['origin_x'] = ox
                    dset.attrs['origin_y'] = oy
                except MemoryError:
                    print(f"[WARN] Patch {i} of frame {frame_idx} 写入失败：内存不足，跳过该patch")
        except Exception as e:
            print(f"[ERROR] 写入 frame_{frame_idx} 失败：{e}")
        with lock:
            global total_processed
            total_processed += 1

# ============== 单帧任务：切片，写入队列 ==================
def process_frame_enqueue(frame_idx, frame, slice_size, overlap_ratio):
    slices = slice_image(
        image=frame,
        slice_height=slice_size,
        slice_width=slice_size,
        overlap_height_ratio=overlap_ratio,
        overlap_width_ratio=overlap_ratio
    )
    patch_images = []
    patch_origins = []
    for s in slices:
        patch_images.append(s["image"])
        ox, oy = s["starting_pixel"]
        patch_origins.append((ox, oy))

    while write_queue.full():
        time.sleep(0.05)
    write_queue.put((frame_idx, patch_images, patch_origins))

# ================= 主函数 ===================
def slice_video_with_write_queue(video_path, hdf5_path, slice_size, overlap_ratio, max_workers=8):
    global total_submitted, active_futures, monitor_stop

    # 启动解码线程
    decode_thread = threading.Thread(target=decode_thread_func, args=(video_path,))
    decode_thread.start()

    # 启动监控线程
    monitor_thread_obj = threading.Thread(target=monitor_thread)
    monitor_thread_obj.start()

    # 启动写入线程
    h5_file = h5py.File(hdf5_path, 'w')
    writer_thread = threading.Thread(target=write_worker, args=(h5_file,))
    writer_thread.start()

    executor = ThreadPoolExecutor(max_workers=max_workers)
    frame_idx = 0

    while True:
        frame = frame_queue.get()
        if frame is None:
            break

        # 控制提交速率
        while True:
            with lock:
                active_futures = [f for f in active_futures if not f.done()]
                if len(active_futures) <= max_workers * 4 and not write_queue.full():
                    break
            time.sleep(0.05)

        future = executor.submit(
            process_frame_enqueue,
            frame_idx, frame, slice_size, overlap_ratio
        )
        with lock:
            active_futures.append(future)
            total_submitted += 1
        frame_idx += 1

        # 清理已完成的任务
        with lock:
            active_futures = [f for f in active_futures if not f.done()]

    # 等待所有线程完成
    executor.shutdown(wait=True)
    write_queue.put(None)  # 通知写线程退出
    writer_thread.join()
    h5_file.close()

    monitor_stop = True
    monitor_thread_obj.join()
    decode_thread.join()

    print(f"[✅ DONE] 共处理 {frame_idx} 帧，切片结果已保存至 {hdf5_path}")

# ================= 入口 ===================
if __name__ == '__main__':
    OVERLAP_RATIO = 0.2
    MAX_WORKERS = 8

    print("[STEP 1] 开始并行切片+写入...")
    slice_video_with_write_queue(
        video_path=INPUT_VIDEO_PATH,
        hdf5_path=HDF5_PATH,
        slice_size=SLICE_SIZE,
        overlap_ratio=OVERLAP_RATIO,
        max_workers=MAX_WORKERS
    )
